{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d904dee",
   "metadata": {},
   "source": [
    "# Example 2: Speeding up\n",
    "\n",
    "Major concerns about KANs are their slow running speed and huge memory. This is mainly due to the naive implementation of the first version. We have done a few efficiency updates in the new release. \n",
    "\n",
    "* We update the spline evaluation method, inspired from the efficientKAN repo.\n",
    "* We provide a method to speed up training, simply call model = model.speed(). In this speed mode, the symbolic front is skipped (which will save computation time), and intermediate activations are not saved (which save memory)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc709a73",
   "metadata": {},
   "source": [
    "### Below we compare the normal mode and the speed mode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94056ef6",
   "metadata": {},
   "source": [
    "The Normal mode without speeding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a59179d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Directory already exists: ./model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train loss: 1.45e+01 | test loss: 2.91e+01 | reg: 6.95e+03 : 100%|██| 10/10 [00:06<00:00,  1.66it/s]\n"
     ]
    }
   ],
   "source": [
    "from kan import *\n",
    "\n",
    "seed = 1\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# initialize KAN with G=3\n",
    "model = KAN(width=[2,30,30,1], grid=3, k=3, seed=1)\n",
    "\n",
    "# create dataset\n",
    "f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)\n",
    "dataset = create_dataset(f, n_var=2)\n",
    "model.fit(dataset, opt=\"Adam\", steps=10, update_grid=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffe10f34",
   "metadata": {},
   "source": [
    "The speed mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1233ed39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Directory already exists: ./model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train loss: 1.53e+01 | test loss: 2.52e+01 | reg: 0.00e+00 : 100%|██| 10/10 [00:00<00:00, 10.03it/s]\n"
     ]
    }
   ],
   "source": [
    "from kan import *\n",
    "\n",
    "seed = 1\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# initialize KAN with G=3\n",
    "model = KAN(width=[2,30,30,1], grid=3, k=3, seed=1)\n",
    "model = model.speed()\n",
    "\n",
    "# create dataset\n",
    "f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)\n",
    "dataset = create_dataset(f, n_var=2)\n",
    "model.fit(dataset, opt=\"Adam\", steps=10, update_grid=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "414ec95d",
   "metadata": {},
   "source": [
    "However, the speed mode does not save intermediate activations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c521e5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.acts_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1aa0659e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[ 7.6856, 16.2691],\n",
       "         [15.3356,  4.7036],\n",
       "         [ 1.7350,  0.5762],\n",
       "         [13.7567,  0.6547],\n",
       "         [ 2.0027,  1.1051],\n",
       "         [ 1.3028, 28.8967],\n",
       "         [ 3.1434,  0.9702],\n",
       "         [20.5310,  0.9579],\n",
       "         [10.5522,  3.1435],\n",
       "         [ 1.8792,  2.7931],\n",
       "         [27.2825,  5.3805],\n",
       "         [ 5.1843,  0.9814],\n",
       "         [ 5.2057,  5.5816],\n",
       "         [ 9.4054,  0.3300],\n",
       "         [ 8.7830,  1.9473],\n",
       "         [ 1.6670,  5.7267],\n",
       "         [ 6.4220, 28.1405],\n",
       "         [18.4623, 10.7246],\n",
       "         [33.4021,  0.9486],\n",
       "         [ 1.6931,  3.2448],\n",
       "         [ 7.3774, 16.6039],\n",
       "         [ 1.5200,  1.4180],\n",
       "         [ 6.8054,  0.9019],\n",
       "         [ 1.9008,  1.0136],\n",
       "         [ 0.4024,  4.3498],\n",
       "         [ 4.3653,  5.3940],\n",
       "         [ 4.2989,  3.5657],\n",
       "         [ 4.5522,  0.4301],\n",
       "         [ 0.5495,  0.2769],\n",
       "         [24.9288,  4.5964]]),\n",
       " tensor([[4.4354e+00, 5.6637e-02, 4.2377e+00, 8.8434e-01, 3.7175e+00, 8.9901e-01,\n",
       "          2.0012e+00, 7.9254e-01, 1.5717e+00, 8.3965e-12, 9.4111e-01, 2.7945e+00,\n",
       "          2.3130e-01, 1.0786e+00, 2.0924e+00, 1.8385e-06, 7.8853e-01, 1.8605e+00,\n",
       "          1.0704e+00, 3.3280e-08, 9.1237e-01, 1.3574e-14, 7.1152e-14, 9.1538e-01,\n",
       "          1.8903e+00, 2.8360e+00, 1.7248e-09, 2.8683e+00, 1.6309e+00, 2.7791e-01],\n",
       "         [2.7741e+00, 1.7333e-01, 4.0277e+00, 6.1746e-01, 2.4397e+00, 8.5001e-01,\n",
       "          9.8325e-01, 1.2499e+00, 3.6122e-01, 1.0975e-11, 1.0640e+00, 2.7793e+00,\n",
       "          1.5701e-01, 2.3363e+00, 2.7972e+00, 2.0593e-06, 1.0298e-01, 2.1000e+00,\n",
       "          1.0845e+00, 3.5262e-08, 1.0165e+00, 1.3794e-14, 1.2865e-13, 2.2997e+00,\n",
       "          1.6426e+00, 3.1206e+00, 8.2411e-10, 1.7321e+00, 1.3940e+00, 2.0221e-01],\n",
       "         [2.4054e+00, 9.5195e-02, 4.0988e+00, 6.9150e-01, 3.1943e+00, 9.6773e-01,\n",
       "          2.2461e+00, 1.4729e+00, 1.1388e+00, 1.2414e-11, 9.4799e-01, 2.6189e+00,\n",
       "          2.8233e-01, 2.0001e+00, 2.3587e+00, 2.0759e-06, 4.4866e-01, 2.0386e+00,\n",
       "          2.0690e+00, 3.4758e-08, 1.5009e+00, 1.3446e-14, 1.1570e-13, 5.5674e+00,\n",
       "          3.6710e+00, 2.6040e+00, 8.2651e-10, 2.3744e+00, 1.1822e+00, 3.1520e-01],\n",
       "         [5.2983e+00, 7.3604e-02, 2.9882e+00, 4.7902e-01, 1.6752e-01, 4.7246e-01,\n",
       "          5.3422e-01, 1.0296e-01, 9.3000e-01, 1.8681e-11, 1.4799e+00, 3.1778e-01,\n",
       "          2.4206e-02, 1.2860e+00, 4.8595e+00, 1.3563e-06, 2.8412e-02, 2.2736e+00,\n",
       "          2.5832e-01, 4.4527e-08, 7.3574e-01, 8.9308e-17, 1.1696e-13, 1.8853e+00,\n",
       "          2.9812e+00, 2.2684e+00, 1.3208e-09, 2.3178e+00, 7.5722e+00, 3.8626e-01],\n",
       "         [2.4889e+00, 2.0212e-02, 5.7652e-01, 4.3576e-01, 2.0011e+00, 5.3115e-01,\n",
       "          2.2109e+00, 6.4719e-01, 3.8911e-02, 1.2620e-12, 6.4222e-01, 4.5835e-01,\n",
       "          2.3331e-02, 7.6864e-01, 2.1965e+00, 7.5926e-07, 8.3256e-01, 3.5998e+00,\n",
       "          1.2404e+00, 2.4087e-08, 4.6261e-01, 7.1813e-15, 3.8952e-14, 1.0252e+00,\n",
       "          3.1197e-01, 4.3089e-01, 5.9730e-10, 1.1181e+00, 4.6878e+00, 7.3552e-02],\n",
       "         [2.6117e+00, 1.0827e-01, 3.9609e+00, 9.2829e-01, 2.2001e+00, 7.6530e-01,\n",
       "          3.3549e+00, 1.0102e+00, 1.6966e+00, 7.7765e-12, 7.7808e-01, 3.3234e+00,\n",
       "          2.2734e-01, 1.8369e+00, 2.6977e+00, 1.9341e-06, 3.9897e-01, 1.8150e+00,\n",
       "          1.4196e+00, 3.1108e-08, 1.1459e+00, 1.0992e-14, 1.1278e-13, 5.5055e+00,\n",
       "          1.8089e+00, 3.0406e+00, 9.9519e-10, 3.1057e+00, 4.9756e-01, 3.5589e-01],\n",
       "         [2.4007e+00, 1.4260e-01, 5.2718e+00, 1.0715e+00, 2.5959e+00, 1.6713e+00,\n",
       "          8.7972e-01, 9.2909e-01, 1.0689e+00, 1.6891e-11, 1.3322e+00, 9.7792e-01,\n",
       "          3.5284e-01, 5.8153e-01, 5.2583e+00, 1.3193e-06, 8.9683e-01, 3.7310e+00,\n",
       "          1.8045e+00, 2.5081e-08, 1.5504e+00, 1.7020e-14, 9.6472e-14, 1.4833e+00,\n",
       "          2.4796e+00, 3.7008e+00, 9.0683e-10, 2.5865e+00, 1.4228e+01, 4.0819e-01],\n",
       "         [4.5005e+00, 1.1403e-01, 2.7783e+00, 4.9161e-01, 4.9396e+00, 3.6857e-01,\n",
       "          2.8450e+00, 1.8193e+00, 9.4404e-01, 2.4453e-12, 3.5365e-01, 3.1969e-01,\n",
       "          2.7763e-01, 1.0288e+00, 3.6257e-01, 8.2115e-07, 2.7806e-01, 1.1331e+00,\n",
       "          7.4675e-01, 5.9031e-09, 4.8439e-01, 7.1289e-15, 2.9170e-14, 2.9636e+00,\n",
       "          1.7487e+00, 1.4293e+00, 3.4924e-10, 6.2544e-01, 2.0350e+00, 1.5112e-01],\n",
       "         [2.7216e+00, 1.6853e-01, 4.2689e+00, 7.1123e-01, 1.2633e+00, 9.7850e-01,\n",
       "          1.1939e+00, 5.9262e-01, 1.7410e+00, 8.4416e-12, 1.0811e+00, 3.7561e+00,\n",
       "          2.3250e-01, 1.7598e+00, 3.0913e+00, 1.9736e-06, 3.1038e-01, 2.0196e+00,\n",
       "          1.9533e+00, 3.4704e-08, 1.1332e+00, 1.1547e-14, 1.4854e-13, 6.6080e+00,\n",
       "          1.8311e+00, 3.2493e+00, 8.9923e-10, 3.0517e+00, 2.0265e+00, 2.4044e-01],\n",
       "         [3.0166e+00, 7.2694e-03, 2.4098e+00, 6.4514e-01, 5.3105e+00, 4.3760e-01,\n",
       "          4.2899e-01, 1.6486e+00, 2.8066e-01, 1.6229e-11, 5.9427e-01, 2.1634e+00,\n",
       "          2.0958e-01, 6.1915e-01, 2.2383e+00, 9.8935e-07, 2.6814e-01, 1.4322e+00,\n",
       "          1.5998e+00, 3.6335e-08, 1.7243e+00, 1.1170e-14, 3.4551e-14, 1.7256e+00,\n",
       "          1.2273e+00, 2.8237e+00, 2.0222e-10, 9.3636e-01, 8.0074e-01, 1.8898e-01],\n",
       "         [1.6953e+00, 2.0251e-02, 5.3069e+00, 3.0458e-01, 4.4163e+00, 4.9125e-01,\n",
       "          4.9296e-01, 1.8665e+00, 2.6483e-01, 3.3367e-12, 3.9696e-01, 3.3676e-01,\n",
       "          1.6086e-01, 8.4875e-01, 5.2545e+00, 7.3239e-07, 1.6651e-01, 9.3421e-01,\n",
       "          3.8561e-01, 1.1715e-08, 3.9553e-01, 1.2336e-14, 4.8541e-14, 2.0979e+00,\n",
       "          2.7597e+00, 3.1390e+00, 2.3438e-10, 2.2403e+00, 8.1404e+00, 6.3760e-02],\n",
       "         [2.8035e+00, 1.6375e-01, 2.7706e+00, 7.1630e-01, 4.4769e+00, 9.4805e-01,\n",
       "          2.7471e+00, 1.7877e+00, 3.1662e-01, 8.2097e-12, 1.0638e+00, 5.4979e-01,\n",
       "          2.4679e-01, 1.8376e+00, 2.8678e+00, 1.7142e-06, 5.9221e-01, 2.1449e+00,\n",
       "          2.0114e+00, 4.0508e-08, 1.0354e+00, 1.0502e-14, 1.2674e-13, 7.4354e+00,\n",
       "          2.7504e+00, 3.0467e+00, 7.3362e-10, 1.7270e+00, 6.5320e-01, 3.0570e-01],\n",
       "         [2.2427e+00, 1.7672e-01, 2.3974e+00, 9.0831e-01, 4.1116e+00, 8.0689e-01,\n",
       "          2.5742e+00, 1.8215e+00, 1.0106e+00, 1.0420e-11, 7.8891e-01, 1.9371e+00,\n",
       "          2.2681e-01, 1.2273e+00, 2.2120e+00, 1.8724e-06, 4.1771e-01, 1.8459e+00,\n",
       "          1.3102e+00, 3.2487e-08, 1.0789e+00, 1.5243e-14, 1.1819e-13, 8.3267e+00,\n",
       "          2.1390e+00, 2.6367e+00, 1.0522e-09, 1.8955e+00, 2.9084e+00, 2.8246e-01],\n",
       "         [3.8211e+00, 4.3302e-02, 1.5532e+00, 6.5081e-01, 2.2713e+00, 4.2253e-01,\n",
       "          2.9351e+00, 1.9489e+00, 2.9210e-01, 3.4494e-12, 2.9242e-01, 4.3019e-01,\n",
       "          1.2177e-01, 4.9818e-01, 4.4591e+00, 1.0103e-06, 5.2056e-01, 6.5539e-02,\n",
       "          2.1910e-01, 4.6610e-09, 5.3929e-01, 4.7059e-15, 7.8409e-14, 5.6421e-01,\n",
       "          3.7233e+00, 9.2992e-01, 1.6597e-10, 3.5818e+00, 1.5577e+00, 1.7590e-01],\n",
       "         [2.9046e+00, 9.5698e-02, 2.6694e+00, 6.4469e-01, 4.3034e+00, 9.9968e-01,\n",
       "          2.7034e+00, 1.6186e+00, 1.0375e+00, 1.2487e-11, 1.0187e+00, 3.4804e+00,\n",
       "          3.8877e-01, 9.3618e-01, 2.8343e+00, 1.8381e-06, 4.3643e-01, 2.5044e+00,\n",
       "          1.1255e+00, 3.6622e-08, 1.0358e+00, 9.3915e-15, 1.3382e-13, 3.9234e+00,\n",
       "          2.4406e+00, 2.5703e+00, 9.8101e-10, 1.8558e+00, 1.7584e+00, 2.6487e-01],\n",
       "         [8.2608e-01, 3.9826e-02, 4.9642e+00, 5.9423e-01, 2.9735e+00, 2.9637e-01,\n",
       "          1.1755e+00, 9.5304e-01, 2.2630e-01, 1.9392e-11, 1.4972e-01, 4.1555e+00,\n",
       "          3.3123e-01, 1.2222e+00, 1.2493e+00, 1.1486e-06, 2.4019e-01, 5.5385e-01,\n",
       "          1.2639e+00, 1.7620e-08, 1.8553e+00, 2.8778e-14, 6.3127e-14, 6.0493e+00,\n",
       "          1.6828e+00, 1.6001e+00, 1.8324e-09, 7.2790e-01, 1.2600e+00, 1.6318e-01],\n",
       "         [4.0348e+00, 3.4401e-02, 3.1383e+00, 5.4052e-01, 4.8957e+00, 4.5683e-01,\n",
       "          5.5182e-01, 1.7985e+00, 1.0728e+00, 4.8572e-12, 4.9333e-01, 2.2659e+00,\n",
       "          2.9328e-01, 1.8455e+00, 1.5274e+00, 4.3244e-07, 2.4830e-01, 1.4385e+00,\n",
       "          6.4562e-01, 1.3089e-08, 3.0059e-01, 8.7514e-15, 1.0119e-14, 3.0255e+00,\n",
       "          2.9596e+00, 1.8410e+00, 7.2461e-10, 6.7798e-01, 6.9710e-01, 7.4898e-02],\n",
       "         [2.3273e+00, 1.0593e-01, 2.6666e+00, 7.0182e-01, 1.9286e+00, 1.0449e+00,\n",
       "          2.2420e+00, 6.7089e-01, 1.3388e+00, 1.4481e-11, 5.4829e-01, 7.3503e-01,\n",
       "          9.7126e-02, 2.1228e+00, 3.3200e+00, 2.0983e-06, 3.3594e-01, 2.3478e+00,\n",
       "          1.7114e+00, 3.6809e-08, 9.2607e-01, 1.4859e-14, 1.2494e-13, 5.4480e+00,\n",
       "          2.4350e+00, 3.1165e+00, 1.0926e-09, 2.3066e+00, 1.9802e+00, 3.2123e-01],\n",
       "         [3.5227e+00, 1.1639e-01, 2.9583e+00, 5.0209e-01, 4.8553e+00, 5.0531e-01,\n",
       "          3.7937e+00, 1.4786e+00, 1.9596e+00, 4.3690e-12, 6.9199e-01, 1.9448e+00,\n",
       "          1.5100e-01, 1.3992e+00, 1.3800e+00, 6.1192e-07, 2.1646e-01, 1.6881e+00,\n",
       "          8.3899e-01, 1.0043e-08, 4.8695e-01, 9.1885e-15, 2.8802e-14, 3.4677e+00,\n",
       "          2.1111e+00, 1.6565e+00, 6.7751e-10, 1.3383e+00, 5.6327e+00, 1.5579e-01],\n",
       "         [2.6507e-01, 9.3874e-02, 3.7191e+00, 4.9744e-01, 3.6279e+00, 1.7210e-01,\n",
       "          5.5143e+00, 9.8101e-01, 1.2910e-01, 9.5707e-12, 1.9589e-01, 9.8307e-01,\n",
       "          3.6569e-01, 2.2733e+00, 6.4200e-02, 1.0380e-06, 6.1132e-01, 4.4997e-01,\n",
       "          1.0535e+00, 4.2998e-09, 8.5013e-01, 2.4626e-14, 6.8968e-14, 4.9795e+00,\n",
       "          1.6740e+00, 1.8986e+00, 1.9389e-09, 3.5265e+00, 3.2236e+00, 7.8258e-02],\n",
       "         [2.6591e+00, 7.9823e-02, 3.3652e+00, 8.2330e-01, 4.3604e+00, 8.3722e-01,\n",
       "          4.1932e+00, 1.5685e+00, 1.5388e+00, 1.2528e-11, 8.2972e-01, 2.9508e+00,\n",
       "          3.8028e-01, 2.0163e+00, 3.1038e+00, 1.0608e-06, 5.5961e-01, 2.1320e+00,\n",
       "          1.9303e+00, 3.6078e-08, 1.0092e+00, 1.7196e-14, 1.1954e-13, 8.1442e+00,\n",
       "          2.6229e+00, 2.9574e+00, 1.0425e-09, 3.0335e+00, 1.5076e+01, 2.3407e-01],\n",
       "         [2.6346e+00, 1.0826e-01, 2.3256e+00, 9.8563e-01, 2.0622e+00, 9.6685e-01,\n",
       "          6.1216e-01, 5.1432e-01, 1.1011e+00, 8.4878e-12, 9.0861e-01, 1.4643e+00,\n",
       "          2.2213e-01, 1.7886e+00, 3.5816e+00, 2.0015e-06, 3.5627e-01, 2.7970e+00,\n",
       "          1.4387e+00, 3.5869e-08, 8.4443e-01, 1.2872e-14, 1.2780e-13, 5.5024e+00,\n",
       "          3.9389e+00, 2.7262e+00, 9.9537e-10, 1.9653e+00, 2.5938e+00, 2.3006e-01],\n",
       "         [4.2678e+00, 1.9903e-02, 3.8227e+00, 6.5938e-01, 5.1752e+00, 3.6617e-01,\n",
       "          6.0096e-01, 1.4982e+00, 1.9169e-01, 3.5221e-12, 6.0433e-01, 1.9519e+00,\n",
       "          2.4719e-01, 7.8279e-01, 1.5779e+00, 9.7484e-07, 1.5838e-01, 1.8141e+00,\n",
       "          1.8866e+00, 2.7451e-08, 1.6971e+00, 5.8982e-15, 4.5401e-14, 1.6402e+00,\n",
       "          1.5726e+00, 2.7363e+00, 5.7048e-10, 3.0334e+00, 5.7986e+00, 6.6437e-02],\n",
       "         [4.1412e+00, 1.4827e-01, 6.5206e+00, 1.1945e+00, 3.8590e+00, 1.6201e+00,\n",
       "          3.2736e+00, 1.6046e+00, 1.9894e+00, 1.7346e-11, 1.3626e+00, 3.9205e+00,\n",
       "          1.6820e-01, 2.1799e+00, 7.2165e-01, 2.0063e-06, 8.0039e-01, 3.0162e+00,\n",
       "          1.6942e+00, 1.3580e-09, 1.2274e+00, 2.5733e-14, 1.4663e-13, 5.4378e+00,\n",
       "          3.3899e+00, 3.3364e+00, 1.7597e-09, 1.8485e+00, 1.8154e+00, 4.0665e-01],\n",
       "         [1.7175e+00, 9.6331e-02, 4.2227e+00, 5.7334e-01, 1.6949e+00, 1.0976e+00,\n",
       "          2.8887e+00, 1.0743e+00, 7.4266e-01, 8.8706e-12, 7.6744e-01, 7.8671e-01,\n",
       "          2.5058e-01, 1.9733e+00, 2.5570e+00, 1.9162e-06, 4.6122e-01, 2.3221e+00,\n",
       "          1.5756e+00, 3.2919e-08, 9.5216e-01, 1.4645e-14, 1.3467e-13, 4.6060e+00,\n",
       "          1.9789e+00, 3.2284e+00, 8.3445e-10, 3.0383e+00, 1.8058e+00, 2.2853e-01],\n",
       "         [2.3204e+00, 1.6079e-01, 2.7723e+00, 6.1717e-01, 2.4717e+00, 7.9156e-01,\n",
       "          7.0750e-01, 1.0640e+00, 3.3363e-01, 1.1368e-11, 9.3369e-01, 5.7124e-01,\n",
       "          1.1679e-01, 1.8633e+00, 2.7119e+00, 1.8430e-06, 3.3116e-01, 2.0191e+00,\n",
       "          2.3322e+00, 3.1797e-08, 1.0036e+00, 1.4443e-14, 1.2325e-13, 5.0863e+00,\n",
       "          2.7657e+00, 2.8643e+00, 1.1279e-09, 2.7509e+00, 2.0778e+00, 2.2715e-01],\n",
       "         [4.6174e+00, 4.5653e-02, 4.3581e+00, 3.4115e-01, 3.7316e+00, 6.1426e-01,\n",
       "          1.0902e+00, 1.8812e+00, 1.2897e+00, 1.0610e-11, 5.2964e-01, 1.1701e+00,\n",
       "          9.6843e-02, 1.2852e+00, 3.5282e+00, 1.2262e-06, 2.7302e-01, 8.1845e-01,\n",
       "          9.4639e-01, 1.5690e-08, 5.3619e-01, 4.6107e-15, 1.3429e-13, 2.3928e+00,\n",
       "          4.2985e+00, 6.9696e-01, 5.2320e-10, 1.5178e+00, 1.6299e+00, 2.9501e-01],\n",
       "         [4.5484e+00, 8.3778e-02, 4.0664e+00, 1.2043e+00, 3.1747e+00, 2.1460e+00,\n",
       "          2.4384e+00, 1.0987e+00, 3.9380e-01, 1.9647e-11, 2.0224e+00, 3.7419e-01,\n",
       "          3.0129e-02, 1.9915e-01, 4.4007e+00, 1.3485e-06, 1.2234e+00, 5.2071e+00,\n",
       "          1.1689e+00, 3.9132e-08, 7.5287e-01, 2.6759e-14, 5.8267e-14, 4.0041e+00,\n",
       "          6.2831e-01, 1.0457e+00, 1.4252e-09, 1.8536e+00, 8.2997e-01, 4.2599e-01],\n",
       "         [4.9122e+00, 1.4955e-02, 5.1885e+00, 5.4270e-01, 2.5415e+00, 4.8918e-01,\n",
       "          2.9930e+00, 1.8103e+00, 4.5803e-01, 2.7637e-12, 1.7019e-01, 9.7642e-01,\n",
       "          1.8464e-01, 5.7884e-01, 4.5927e+00, 9.7754e-07, 2.2764e-01, 7.6643e-01,\n",
       "          3.6063e-01, 1.1229e-08, 3.1307e-01, 3.3746e-15, 7.7355e-14, 1.9151e+00,\n",
       "          4.0543e+00, 1.0374e+00, 3.7742e-10, 1.4498e+00, 3.4803e+00, 1.4359e-01],\n",
       "         [4.4921e+00, 2.6448e-02, 2.4570e+00, 6.0745e-01, 3.1722e+00, 3.1898e-01,\n",
       "          6.5110e-01, 1.4345e+00, 2.2791e-01, 1.4314e-12, 1.8415e-01, 1.7340e-01,\n",
       "          7.6617e-02, 2.1373e+00, 1.6081e+00, 7.4545e-07, 2.1735e-01, 4.0267e-01,\n",
       "          2.8377e-01, 1.2359e-08, 4.7323e-01, 5.2978e-15, 6.7723e-14, 1.9885e+00,\n",
       "          4.3402e+00, 9.6498e-01, 3.8371e-10, 9.8866e-01, 2.1566e+00, 1.9733e-01]]),\n",
       " tensor([[8.0486e-23, 0.0000e+00, 6.6755e-20, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00, 1.6415e-10, 0.0000e+00, 2.9900e-02, 2.8350e-16, 1.0642e-20,\n",
       "          2.1081e-22, 2.0802e-01, 8.2120e-14, 0.0000e+00, 6.9611e-08, 0.0000e+00,\n",
       "          2.0685e-02, 0.0000e+00, 7.1422e-22, 0.0000e+00, 1.9870e-02, 1.4698e-02,\n",
       "          0.0000e+00, 1.3654e-33, 4.5392e-02, 0.0000e+00, 1.3617e-01, 1.4280e-13]])]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save_plot_data = True\n",
    "model.get_act(dataset)\n",
    "model.acts_scale"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
